import os
import sys
import csv
import time
import json
import torch
from core.deploy import MiniCPMVDeploy
from inference.image_inference import ImageInferencer
from utils.batch_processor import process_directory, get_image_files_in_directory

class BatchImageInference:
    """
    MiniCPM-V 批量图片推理工具，基于多任务推理服务器实现模型一次加载多次使用
    """
    def __init__(self, config_file=None):
        # 初始化部署类（不立即加载模型）
        self.deploy = MiniCPMVDeploy(config_file)
        self.model_loaded = False
        
    def load_model(self):
        """加载模型（只需执行一次）"""
        if not self.model_loaded:
            print("\n===== 开始加载模型 =====")
            start_time = time.time()
            if self.deploy.load_model():
                self.model_loaded = True
                load_time = time.time() - start_time
                print(f"模型加载成功，耗时: {load_time:.2f}秒")
                print("模型已加载至内存，可进行多次推理任务")
            else:
                print("模型加载失败，请检查模型配置和路径")
                return False
        else:
            print("模型已经加载，无需重复加载")
        return True
        
    def batch_process(self, folder_path, output_csv, prompt_template=None, max_workers=None, recursive=True, no_save=False):
        """
        批量处理图片文件夹，优化版本，实现模型一次加载多次使用并提高性能
        
        Args:
            folder_path: 图片文件夹路径
            output_csv: 输出CSV文件路径
            prompt_template: 提示词模板，支持变量替换
            max_workers: 最大工作线程数
            recursive: 是否递归搜索子目录
            no_save: 是否不保存单个结果文件
        
        Returns:
            是否处理成功
        """
        if not self.model_loaded:
            print("错误: 模型尚未加载，请先加载模型")
            return False
        
        # 检查文件夹是否存在
        if not os.path.exists(folder_path):
            print(f"错误: 文件夹 {folder_path} 不存在")
            return False
        
        # 如果没有提供提示词，使用默认提示词
        if not prompt_template:
            prompt_template = "请详细描述这张图片"
        
        print(f"\n===== 开始批量图片推理 =====")
        print(f"图片目录: {folder_path}")
        print(f"提示词模板: {prompt_template}")
        print(f"是否递归搜索: {recursive}")
        if max_workers is not None:
            print(f"最大工作线程数: {max_workers}")
        if output_csv:
            print(f"CSV导出路径: {output_csv}")
        print(f"是否保存单个结果文件: {not no_save}")
        
        # 获取图像文件列表
        image_files = get_image_files_in_directory(folder_path, recursive)
        if not image_files:
            print(f"在目录{folder_path}中未找到图像文件")
            return False
        
        print(f"找到{len(image_files)}个图像文件待处理")
        
        # 创建推理器实例
        inferencer = ImageInferencer(self.deploy)
        
        # 准备输出目录
        if output_csv:
            output_dir = os.path.dirname(output_csv)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
        
        # 准备CSV文件
        csv_file = None
        csv_writer = None
        if output_csv:
            csv_file = open(output_csv, 'w', newline='', encoding='utf-8')
            fieldnames = ['filename', 'file_path', 'result', 'time_cost']
            csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
            csv_writer.writeheader()
        
        # 批量处理，使用串行模式以避免内存问题
        results = []
        total_files = len(image_files)
        success_count = 0
        total_time_cost = 0
        
        for i, image_path in enumerate(image_files):
            print(f"\n[{i+1}/{total_files}] 处理图像: {os.path.basename(image_path)}")
            
            # 打印当前内存使用情况
            if torch.cuda.is_available():
                used_memory = torch.cuda.memory_allocated() / 1024**3
                print(f"当前GPU内存使用: {used_memory:.2f} GB")
            
            # 准备提示词
            filename = os.path.basename(image_path)
            prompt = prompt_template.replace('{file_name}', filename) if prompt_template else None
            
            # 执行推理
            start_time = time.time()
            result = inferencer.infer(image_path, prompt, not no_save)
            end_time = time.time()
            time_cost = end_time - start_time
            
            # 记录结果
            results.append((image_path, result, time_cost))
            
            # 更新统计信息
            if result and "处理失败" not in str(result):
                success_count += 1
            total_time_cost += time_cost
            
            # 写入CSV
            if csv_writer:
                row = {
                    'filename': os.path.basename(image_path),
                    'file_path': image_path,
                    'result': result,
                    'time_cost': f"{time_cost:.2f}"
                }
                csv_writer.writerow(row)
            
            # 清理内存
            self._clean_memory()
        
        # 关闭CSV文件
        if csv_file:
            csv_file.close()
            print(f"结果已导出到CSV文件: {output_csv}")
        
        # 统计信息
        failed_count = total_files - success_count
        print(f"\n批量图片推理完成，耗时: {total_time_cost:.2f}秒")
        print(f"统计信息:")
        print(f"- 总图片数: {total_files}")
        print(f"- 成功处理: {success_count}")
        print(f"- 失败处理: {failed_count}")
        print(f"- 总耗时: {total_time_cost:.2f}秒")
        print(f"- 平均耗时: {total_time_cost/total_files:.2f}秒/张")
        
        return True
        
    def _clean_memory(self):
        """清理内存，防止内存泄漏"""
        # 清理PyTorch缓存
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        # 强制Python进行垃圾回收
        import gc
        gc.collect()
        
    def unload_model(self):
        """卸载模型，释放内存"""
        if self.model_loaded:
            self.deploy.unload_model()
            self.model_loaded = False
            print("模型已卸载，内存已释放")
        else:
            print("模型尚未加载，无需卸载")

def main():
    """主函数"""
    print("===== MiniCPM-V 批量图像处理工具 =====")
    print("基于多任务推理服务器实现模型一次加载，多次使用")
    
    # 创建批量推理实例
    batch_infer = BatchImageInference()
    
    # 加载模型（只需加载一次）
    print("\n正在加载模型...")
    if not batch_infer.load_model():
        print("模型加载失败，程序退出")
        return
    
    # 获取用户输入
    print("\n请输入图像文件夹路径: ")
    folder_path = input().strip()
    if not folder_path:
        print("错误: 图像文件夹路径不能为空")
        return
    
    # 验证文件夹路径
    if not os.path.isdir(folder_path):
        print(f"错误: {folder_path} 不是有效的文件夹路径")
        return
    
    # 获取输出CSV文件路径
    print("请输入输出CSV文件路径 (默认: outputs/batch_image_results.csv): ")
    output_csv = input().strip()
    if not output_csv:
        output_csv = "outputs/batch_image_results.csv"
    
    # 获取提示词模板
    print("请输入推理提示词模板 (可选，支持变量替换): ")
    prompt_template = input().strip()
    
    # 获取并行线程数
    print("请输入并行处理线程数 (可选，留空自动调整): ")
    max_workers_input = input().strip()
    max_workers = None
    if max_workers_input and max_workers_input.isdigit():
        max_workers = int(max_workers_input)
    
    # 是否递归搜索子目录
    print("是否递归搜索子目录？(y/n，默认: y): ")
    recursive_input = input().strip().lower()
    recursive = True
    if recursive_input == 'n':
        recursive = False
    
    # 是否保存单个结果文件
    print("是否保存单个结果文件？(y/n，默认: y): ")
    save_single_input = input().strip().lower()
    no_save = False
    if save_single_input == 'n':
        no_save = True
    
    # 开始批量处理
    print("\n开始批量处理图像...")
    start_total_time = time.time()
    
    if batch_infer.batch_process(
        folder_path=folder_path,
        output_csv=output_csv,
        prompt_template=prompt_template,
        max_workers=max_workers,
        recursive=recursive,
        no_save=no_save
    ):
        end_total_time = time.time()
        print(f"\n批量处理完成！总耗时: {end_total_time - start_total_time:.2f}秒")
        
        # 询问是否继续处理其他图像或卸载模型
        while True:
            print("\n请选择操作:")
            print("1. 处理其他图像文件夹")
            print("2. 卸载模型并退出")
            choice = input().strip()
            
            if choice == '1':
                # 处理其他图像文件夹
                print("请输入新的图像文件夹路径: ")
                new_folder_path = input().strip()
                if not new_folder_path:
                    print("错误: 图像文件夹路径不能为空")
                    continue
                if not os.path.isdir(new_folder_path):
                    print(f"错误: {new_folder_path} 不是有效的文件夹路径")
                    continue
                
                print("请输入新的输出CSV文件路径: ")
                new_output_csv = input().strip()
                if not new_output_csv:
                    # 生成默认文件名
                    import datetime
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    new_output_csv = f"outputs/batch_image_results_{timestamp}.csv"
                
                print("\n开始处理新的图像文件夹...")
                batch_infer.batch_process(
                    folder_path=new_folder_path,
                    output_csv=new_output_csv,
                    prompt_template=prompt_template,
                    max_workers=max_workers,
                    recursive=recursive,
                    no_save=no_save
                )
            elif choice == '2':
                # 卸载模型并退出
                batch_infer.unload_model()
                print("程序已退出")
                break
            else:
                print("无效的选择，请重新输入")
    else:
        print("批量处理失败")
        batch_infer.unload_model()

if __name__ == "__main__":
    main()